# 导入必要的库
import torch
from transformers import BertModel, BertTokenizer, BertConfig

# -------------------- 代码块 2：加载 BERT 模型和分词器 --------------------
# 通过预训练模型的词典加载分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
# 加载预训练模型的配置
model_config = BertConfig.from_pretrained('bert-base-chinese')
# 通过预训练模型名称和修改后的配置加载 BERT 模型
# attn_implementation="eager" 用于消除 BertSdpaSelfAttention 警告，并显式指定使用手动实现的注意力计算
model = BertModel.from_pretrained('bert-base-chinese', config=model_config, attn_implementation="eager")

# -------------------- 代码块 3：文本编码 --------------------
# 定义要编码的文本列表
text = ['你真的很好看。', '这个牌子的咖啡很好喝。']
encoding_results = list()  # 初始化一个空列表来存储编码结果

# 循环遍历文本列表
for txt in text:
    # 使用分词器对文本进行编码，并进行以下设置：
    # max_length = 10:  将序列的最大长度设置为 10。超过此长度的序列将被截断，小于此长度的序列将被填充。
    # padding = 'max_length': 如果序列长度小于 max_length，则进行 padding 到 max_length。使用 'max_length' 策略会填充到指定的最大长度。
    # truncation = True: 如果序列长度大于 max_length，则进行截断。确保序列长度不超过 max_length。
    encoding_results.append(tokenizer.encode_plus(txt, max_length=10, padding='max_length', truncation=True))

# 打印编码结果，结果包含 input_ids (词的ID), token_type_ids (句子类型ID), attention_mask (注意力掩码)
print(encoding_results)

# -------------------- 代码块 4：将 ID 转换为 Token --------------------
# 循环遍历编码结果列表
for res in encoding_results:
    # 使用分词器的 convert_ids_to_tokens 方法将 input_ids 转换为 token 列表，以便于理解编码后的文本内容
    print(tokenizer.convert_ids_to_tokens(res['input_ids']))

# -------------------- 代码块 5：准备模型输入张量 --------------------
input_ids = list()  # 初始化 input_ids 列表，用于存储所有文本的 input_ids
type_ids = list()  # 初始化 token_type_ids 列表，用于存储所有文本的 token_type_ids
mask_ids = list()  # 初始化 attention_mask 列表，用于存储所有文本的 attention_mask

# 循环遍历编码结果列表
for res in encoding_results:
    input_ids.append(res['input_ids'])  # 将 input_ids 添加到 input_ids 列表
    type_ids.append(res['token_type_ids'])  # 将 token_type_ids 添加到 type_ids 列表
    mask_ids.append(res['attention_mask'])  # 将 attention_mask 添加到 mask_ids 列表

# 将列表转换为 PyTorch 张量，以便输入到 BERT 模型中
input_ids = torch.tensor(input_ids)  # 将 input_ids 列表转换为 tensor
type_ids = torch.tensor(type_ids)  # 将 token_type_ids 列表转换为 tensor
mask_ids = torch.tensor(mask_ids)  # 将 attention_mask 列表转换为 tensor

# -------------------- 代码块 6：模型前向传播和输出 --------------------
# 将输入张量传递给 BERT 模型进行前向传播
# token_type_ids = type_ids:  传入 token_type_ids，用于区分不同的句子段落 (虽然在这个例子中只有一个句子)。对于单句输入，通常所有 token_type_ids 都是 0。
# attention_mask = mask_ids: 传入 attention_mask，用于mask padding 部分，防止 attention 机制关注到 padding token，从而影响模型效果。 1 表示实际 token，0 表示 padding token。
outputs = model(input_ids, token_type_ids=type_ids, attention_mask=mask_ids)

# 打印模型输出的 keys，查看模型输出了哪些内容。BERT 模型的输出是一个 BaseModelOutputWithPoolingAndCrossAttentions 对象，它类似于一个字典。
print(outputs.keys())
# 打印 last_hidden_state 的形状，这是最后一层隐藏层的输出。形状为 (batch_size, sequence_length, hidden_size)。
print(outputs['last_hidden_state'].shape)
# 打印 pooler_output 的形状，这是 [CLS] token 经过 Pooler 层的输出，通常用于句子级别的表示。形状为 (batch_size, hidden_size)。
print(outputs['pooler_output'].shape)
